# -*- coding: utf-8 -*-

# 导入必要的库
import pymysql  # MySQL连接
import sqlite3  # SQLite连接
from config import DB_CONFIG, client, MODEL_NAME, SUPPORTED_DB_TYPES, VISUALIZATION_DIR
import re
import os
import logging
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
import warnings
import traceback
from typing import List, Dict, Tuple, Any, Optional, Union

# 忽略Matplotlib警告
warnings.filterwarnings("ignore", category=UserWarning)

# 配置日志
logging.basicConfig(level=logging.INFO,
                    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


class Tools:
    """
    工具类，提供与数据库交互、SQL生成、数据可视化等功能
    包含多个实用方法处理各种操作
    """
    def __init__(self) -> None:
        """初始化工具类"""
        # 确保可视化目录存在
        if not os.path.exists(VISUALIZATION_DIR):
            os.makedirs(VISUALIZATION_DIR)
            logger.info(f'Created visualization directory: {VISUALIZATION_DIR}')
        
        # 默认数据库类型
        self.db_type = 'mysql'  # 默认使用MySQL

    def set_db_type(self, db_type: str) -> bool:
        """
        设置要使用的数据库类型
        
        Args:
            db_type (str): 数据库类型 ('mysql' 或 'sqlite')
            
        Returns:
            bool: 设置是否成功
        """
        if db_type.lower() in SUPPORTED_DB_TYPES:
            self.db_type = db_type.lower()
            logger.info(f'Set database type to: {self.db_type}')
            return True
        else:
            logger.warning(f'Unsupported database type: {db_type}')
            return False
    
    def get_connection(self, db_name: str = 'test_db'):
        """
        创建并返回一个数据库连接，支持多种数据库类型
        
        Args:
            db_name (str): 数据库名称，默认为'test_db'
            
        Returns:
            connection: 数据库连接对象，如果连接失败则返回None
        """
        try:
            # 根据不同的数据库类型建立连接
            if self.db_type == 'mysql':
                # MySQL连接
                connection_info = {**DB_CONFIG, 'database': db_name}
                logger.info(f'Connecting to MySQL database: {db_name}')
                
                conn = pymysql.connect(**DB_CONFIG, database=db_name)
                
                # 测试连接是否有效
                with conn.cursor() as cursor:
                    cursor.execute('SELECT 1')
                    result = cursor.fetchone()
                    if result:
                        logger.info(f'Successfully connected to MySQL database: {db_name}')
                        return conn
                    else:
                        logger.error(f'MySQL connection test failed for: {db_name}')
                        raise Exception('MySQL connection test failed')
            
            elif self.db_type == 'sqlite':
                # SQLite连接
                db_path = f'{db_name}.db'
                logger.info(f'Connecting to SQLite database: {db_path}')
                
                conn = sqlite3.connect(db_path)
                
                # 测试连接是否有效
                cursor = conn.cursor()
                cursor.execute('SELECT 1')
                result = cursor.fetchone()
                if result:
                    logger.info(f'Successfully connected to SQLite database: {db_path}')
                    return conn
                else:
                    logger.error(f'SQLite connection test failed for: {db_path}')
                    raise Exception('SQLite connection test failed')
            
            else:
                # 不支持的数据库类型
                logger.error(f'Unsupported database type: {self.db_type}')
                raise ValueError(f'Unsupported database type: {self.db_type}')
                
        except Exception as e:
            logger.error(f'Connect Database Error: {e}')
            # 返回None便于调用者判断连接状态
            return None

    def show_databases(self):
        """
        获取所有用户可用的数据库列表，排除系统数据库
        
        Returns:
            list: 数据库名称列表，返回空列表表示出错
        """
        try:
            # 定义系统数据库列表，这些数据库通常不要展示给用户
            sys_db = ['information_schema', 'mysql',
                      'performance_schema', 'sys']
                      
            # 使用自动关闭连接的上下文管理器
            with self.get_connection() as conn:
                with conn.cursor() as cursor:
                    # 执行查询语句并获取结果
                    cursor.execute('SHOW DATABASES')
                    dbs = [db[0] for db in cursor.fetchall() if db[0] not in sys_db]
                    
                    logger.info(f'Retrieved {len(dbs)} databases: {dbs}')
                    return dbs
        except Exception as e:
            logger.error(f'Show Databases Error: {e}')
            return []  # 出错时返回空列表

    def show_tables(self, db_name):
        """
        获取指定数据库中的所有表
        
        Args:
            db_name (str): 要查询的数据库名称
            
        Returns:
            list: 表名列表，如果发生错误则返回空列表
        """
        logger.info(f'Getting tables from database: {db_name}')
        try:
            # 强化参数验证
            if not db_name or not isinstance(db_name, str):
                logger.error('Show Tables Error: Database name is empty or invalid')
                return []
            
            # 建立数据库连接
            conn = self.get_connection(db_name)
            if conn is None:
                logger.error(f'Show Tables Error: Could not connect to database {db_name}')
                return []
                
            try:
                with conn.cursor() as cursor:
                    # 查询当前数据库中的所有表
                    cursor.execute('SHOW TABLES')
                    tables = [table[0] for table in cursor.fetchall()]
                    
                    logger.info(f'Found {len(tables)} tables in {db_name}: {tables}')
                    return tables
            finally:
                # 确保连接关闭，避免连接泄漏
                conn.close()
                
        except Exception as e:
            logger.error(f'Show Tables Error: {e}')
            return []  # 确保在出现异常时返回空列表

    # 查询数据表的信息
    def get_table_describe(self, db_name, table_name):
        try:

            desc_info = f'数据库:{db_name},数据表:{table_name},列描述信息(列名|数据类型|备注信息)\n'
            with self.get_connection(db_name) as conn:
                with conn.cursor() as cursor:
                    cursor.execute(
                        f'show full columns from {db_name}.{table_name}')
                    # # 拉取查询结果==>(('customerNumber', 'int', None, 'NO', 'PRI', None, '', 'select,insert,update,references', '').....)
                    # Field| Type| Collation| Null | Key | Default | Extra | Privileges| Comment |
                    desc_info += '\n'.join(
                        [f'{item[0]}|{item[1]}|{item[-1]}' for item in cursor.fetchall()])
            logger.info(f'Get Table Info:{db_name}.{table_name}')
            return desc_info
        except Exception as e:
            logger.error(f'Get Table Info Error:{e}')


    
    # 根据表的描述信息和用户的需求生成查询的SQL
    def generate_sql(self, prompt):
        init_prompt = '''
你是一名专业的SQL查询工程师,擅长根据自然语言需求生成高质量的SQL查询。
你能根据用户输入的描述，识别出所需的数据库、表名、字段名、查询条件等关键元素,并构建出完整的SQL语句。
你要遵循标准SQL语法,表名必须携带数据库信息,生成可读性强、正确、可直接执行的查询语句,生成的SQL中只能有英文内容。
输出的内容将只包含纯粹的SQL代码,没有任何其他说明性文字。
'''.strip()

        try:
            response = client.chat.completions.create(
                model=MODEL_NAME, 
                messages=[
                    {'role': 'system', 'content': init_prompt}, 
                    {'role': 'user', 'content': prompt}
                ]
            )
            logger.info('Generate SQL')
            return response.choices[0].message.content
        except Exception as e:
            logger.error(f'Generate SQL Error:{e},User Prompt:{prompt}')
            raise e

    # 根据数据和用户需求生成可视化的Python代码
    def generate_code(self, prompt):
        init_prompt = '''
你是一名经验丰富的Python开发工程师,能够根据用户输入的描述,识别出可视化的数据特征,并使用Python的matplotlib库生成漂亮的可视化图表。
必须按照以下要求生成完整的可执行代码:
1. 引入必要的库: matplotlib, numpy, pandas 等
2. 添加 plt.figure() 创建新的画布
3. 设置中文字体 plt.rcParams['font.sans-serif']=['SimHei']
4. 绘制图表并添加标题、标签等
5. 使用绝对路径保存图片到 "visualization" 目录
6. 根据图表内容生成有意义的英文文件名(不要使用中文)
7. 最后一行写 "plt" 而不是 "return plt" 或 "plt.show()"
8. 千万不要在代码中使用任何return语句
9. 如果生成条形图(bar chart)，一定要使用垂直方向的条形图(plt.bar)而不是水平方向的(plt.barh)

输出的内容将只包含完整的Python代码,不要包含任何其他说明性文字。
'''.strip()
        
        # 确保可视化目录存在
        vis_dir = 'visualization'
        if not os.path.exists(vis_dir):
            os.makedirs(vis_dir)
            logger.info(f'创建可视化目录: {vis_dir}')
        
        try:
            # 调用API生成代码
            response = client.chat.completions.create(
                model=MODEL_NAME, 
                messages=[
                    {'role': 'system', 'content': init_prompt}, 
                    {'role': 'user', 'content': prompt}
                ]
            )
            logger.info('Generate Code')
            
            # 获取生成的代码
            generated_code = response.choices[0].message.content
            
            # 移除代码中的所有return语句
            generated_code = re.sub(r'return\s+plt', 'plt', generated_code)
            
            # 确保代码末尾是plt而不是plt.show()
            generated_code = re.sub(r'plt\.show\(\)\s*$', 'plt', generated_code)
            
            # 修正文件路径问题 - 替换Linux风格路径为Windows风格
            generated_code = re.sub(r"'/mnt/data/(.+?\.png)'", r'"visualization/\1"', generated_code)
            # 修正其他可能的绝对路径
            generated_code = re.sub(r"'/(.+?\.png)'", r'"visualization/\1"', generated_code)
            # 确保使用正确的可视化目录
            generated_code = re.sub(r'"(.+?)/(.*?\.png)"', r'"visualization/\2"', generated_code)
            
            # 如果代码最后没有plt语句，添加一个
            if not generated_code.strip().endswith('plt'):
                generated_code = generated_code.rstrip() + "\n\n# 返回plt对象\nplt"
            
            # 打印最终生成的代码用于调试
            logger.info(f'Generated visualization code: {generated_code[:100]}...')
                
            return generated_code
                
        except Exception as e:
            logger.error(f'Generate Code Error: {e}')
            raise e

    # 解析出SQL语句从markdown代码块中
    def extract_sql_statements(self, text):
        try:
            # 尝试使用正则表达式提取SQL代码块
            sql_statements = re.findall(r'```sql\n(.*?)\n```', text, re.DOTALL)
            logger.info('Extract SQL')
            return sql_statements[0] if sql_statements else text
        except Exception as e:
            logger.error(f'Extract SQL Error:{e}')
            return text

    # 解析出code
    def extract_code_statements(self, text):
        try:
            # 尝试使用正则表达式提取Python代码块
            code_statements = re.findall(r'```python\n(.*?)\n```', text, re.DOTALL)
            
            # 如果没找到标准Python代码块，尝试无语言的代码块
            if not code_statements:
                code_statements = re.findall(r'```\n(.*?)\n```', text, re.DOTALL)
            
            # 使用第一个匹配的代码块，如果存在
            if code_statements:
                logger.info('Extract Code')
                return code_statements[0]
            
            # 如果不存在代码块格式，但文本有import语句，视为Python代码
            if 'import' in text:
                return text
                
            logger.warning('Extract Code: 未找到代码块')
            return text
        except Exception as e:
            logger.error(f'Extract Code Error:{e}')
            return text

    # 执行sql --> 查询出的数据封装为DataFrame --> gradio Dataframe
    def execute_sql(self, db_name, sql):
        try:
            with self.get_connection(db_name) as conn:
                with conn.cursor() as cursor:
                    cursor.execute(sql)
                    query_result = cursor.fetchall()
                    columns = [column[0] for column in cursor.description]
                    df = pd.DataFrame(data=query_result, columns=columns)

                    for column in df.columns:
                        try:
                            df[column] = df[column].astype(float)
                        except Exception as e:
                            continue
            logger.info('Execute SQL')
            return df
        except Exception as e:
            logger.error(f'Execute SQL Error:{e}')

    # 获取指定目录下所有文件的绝对路径
    def get_all_files_path(self, directory):
        try:
            # 确保目录存在
            if not os.path.exists(directory):
                os.makedirs(directory)
                logger.info(f'创建目录: {directory}')
                
            files_path = []
            for root_dir, _, files in os.walk(directory):
                for item in files:
                    if item.startswith('.'):
                        continue
                    file_path = os.path.join(root_dir, item)
                    files_path.append(file_path)
            logger.info(f'Get All Files Path: {directory}, 找到 {len(files_path)} 个文件')
            return files_path
        except Exception as e:
            logger.error(f'Get All Files Path Error:{e}')
            return []
